import psycopg2
import threading
import queue
from flask import Flask, render_template, jsonify
import dash
import dash_core_components as dcc
import dash_html_components as html
import concurrent.futures

DB_CONFIG = {
    "dbname": "simulation_db",
    "user": "sim_user",
    "password": "your_password",
    "host": "localhost",
    "port": "5432"
}

MAX_RETRIES = 1
RUNNING = True  # 全局控制标志，支持暂停
TASK_QUEUE = queue.Queue()
THREAD_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=10)

app = Flask(__name__)
dashboard = dash.Dash(__name__, server=app, routes_pathname_prefix='/dashboard/')

def multi_simulate(neut, region, universe, start, name, delay=1):
    with TaskManager() as task_manager:
        task_manager.reset_running_tasks()
        while RUNNING:
            pending_tasks = task_manager.get_pending_tasks(limit=10)
            if not pending_tasks:
                break
            
            for task in pending_tasks:
                TASK_QUEUE.put(task)
                THREAD_POOL.submit(process_task, task, neut, region, universe, delay, task_manager, name)

def process_task(task, neut, region, universe, delay, task_manager, name):
    task_id, alpha, decay = task
    task_manager.update_status(task_id, 'running')
    
    success = run_simulation(alpha, decay, neut, region, universe, delay)
    
    if success:
        task_manager.update_status(task_id, 'completed')
    else:
        task_manager.increment_retries(task_id)
        if task_manager.get_retries(task_id) >= MAX_RETRIES:
            task_manager.update_status(task_id, 'failed')
        else:
            task_manager.update_status(task_id, 'pending')

def run_simulation(alpha, decay, neut, region, universe, delay):
    return True  # 模拟 API 调用

class TaskManager:
    def __enter__(self):
        self.conn = psycopg2.connect(**DB_CONFIG)
        self.lock = threading.Lock()
        self._create_table()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.conn.close()

    def _create_table(self):
        with self.conn.cursor() as cur:
            cur.execute("""
                CREATE TABLE IF NOT EXISTS simulation_tasks (
                    id SERIAL PRIMARY KEY,
                    alpha TEXT NOT NULL,
                    decay INTEGER NOT NULL,
                    status TEXT CHECK( status IN ('pending', 'running', 'failed', 'completed') ) DEFAULT 'pending',
                    retries INTEGER DEFAULT 0,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)
            self.conn.commit()

    def get_pending_tasks(self, limit=10):
        with self.lock, self.conn.cursor() as cur:
            cur.execute("SELECT id, alpha, decay FROM simulation_tasks WHERE status IN ('pending', 'failed') LIMIT %s", (limit,))
            return cur.fetchall()

    def update_status(self, task_id, status):
        with self.lock, self.conn.cursor() as cur:
            cur.execute("UPDATE simulation_tasks SET status = %s WHERE id = %s", (status, task_id))
            self.conn.commit()

    def increment_retries(self, task_id):
        with self.lock, self.conn.cursor() as cur:
            cur.execute("UPDATE simulation_tasks SET retries = retries + 1 WHERE id = %s", (task_id,))
            self.conn.commit()

    def get_retries(self, task_id):
        with self.conn.cursor() as cur:
            cur.execute("SELECT retries FROM simulation_tasks WHERE id = %s", (task_id,))
            return cur.fetchone()[0]

    def reset_running_tasks(self):
        with self.conn.cursor() as cur:
            cur.execute("UPDATE simulation_tasks SET status = 'pending' WHERE status = 'running'")
            self.conn.commit()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/progress')
def progress():
    with TaskManager() as task_manager:
        with task_manager.conn.cursor() as cur:
            cur.execute("SELECT status, COUNT(*) FROM simulation_tasks GROUP BY status")
            data = cur.fetchall()
    return jsonify(dict(data))

dashboard.layout = html.Div([
    html.H1("Simulation Progress"),
    dcc.Interval(id='interval-component', interval=5000, n_intervals=0),
    dcc.Graph(id='progress-chart')
])

if __name__ == "__main__":
    THREAD_POOL.submit(multi_simulate, "SUBINDUSTRY", "USA", "TOP3000", 0, "test_run")
    app.run(debug=True)